#include "stdafx.h"

#include "SkinTracker.h"
#include <algorithm>
using namespace cv;

#define CROP_W 128
#define CROP_H 128

#define Tmax 0.5
#define Tmin 0.15
/*
|===============================================|
| P(s|c)	=	P(c|s) * P(s) / P(c)			|
|===============================================|
| 1/P(c)	=	total_pixel / num_color			|
| P(s)		=	num_skin / total_pixel			|
| P(c|s)	=	num_color_be_skin / num_Skin	|
|===============================================|=====================|
| P(s|c)	=	num_color_be_skin / num_color
|=====================================================================|
*/
SkinTracker::SkinTracker()
{
}


SkinTracker::~SkinTracker()
{
}

void SkinTracker::initialize()
{
	m_Psc = Mat::zeros(256,256,CV_64FC1);
	folderPath = "C:\\Users\\HoangQC\\Desktop\\KinectDB\\TrainingDB\\";
	vector<Mat> vImgRGB,vImgMask;

	//Load & train images
	for (int i=0;i<7;i++)
	{
		string fileRGB = folderPath + "color_" + std::to_string(i+1) + ".bmp";
		string fileMask = folderPath + "color_"+ std::to_string(i+1) + "_mask.bmp";
		Mat imRGB = imread(fileRGB, CV_LOAD_IMAGE_UNCHANGED);
		Mat mask = imread(fileMask, CV_LOAD_IMAGE_GRAYSCALE);
		vImgRGB.push_back(imRGB);
		vImgMask.push_back(mask);
		updateTraining(vImgRGB[i],vImgMask[i],i+1);
	}
}

void SkinTracker::updateTraining(Mat & imgBGR,Mat & imgMask, double count)
{
	Mat imgYUV;
	Mat colorTable			= Mat::zeros(256,256,CV_64FC1);
	Mat color_skin_Table	= Mat::zeros(256,256,CV_64FC1);

	vector<Mat> channels;
	double total_pixel = imgBGR.rows*imgBGR.cols;
	double num_skin = 0;

	cvtColor(imgBGR,imgYUV,CV_BGR2YCrCb);
	split(imgYUV, channels);

	Mat u_channel = channels[1];
	Mat v_channel = channels[2];

	for (int i = 0; i < imgYUV.rows; i++)
	{
		for (int j = 0; j < imgYUV.cols; j++)
		{
			uchar uColor = u_channel.at<uchar>(i,j);
			uchar vColor = v_channel.at<uchar>(i,j);
			colorTable.at<double>(uColor,vColor) += 1;	//	count Color
			if (imgMask.at<uchar>(i,j) > 128)
			{
				color_skin_Table.at<double>(uColor,vColor) += 1;	//	count Color-Skin
			}

		}
	}

	/*Probability Update*/
	assert(count>0);

	double keepingFactor = (count-1)/count;
	double updateFactor = 1/count;
	for (int i = 0; i < 256; i++)
	{
		for (int j = 0; j < 256; j++)
		{
			double currentColor = colorTable.at<double>(i,j);
			double currentColorBeSkin = color_skin_Table.at<double>(i,j);

			m_Psc.at<double>(i,j) = m_Psc.at<double>(i,j)*keepingFactor;
			if (currentColorBeSkin !=0)
			{
				m_Psc.at<double>(i,j) += updateFactor*currentColorBeSkin/currentColor;
			}
		}
	}
}

void SkinTracker::trackColorObject(Mat & imgBGR)
{
	
}

/* Depth contour for seeding */
void SkinTracker::trackFramesBGR(Mat & imgBGR, Mat & imgD, Point2i center, ushort depthThresh)
{
	Mat imgYUV;
	Mat preview = Mat::zeros(CROP_H,CROP_W,CV_8UC1);
	Mat depthMask = Mat::zeros(CROP_H,CROP_W,CV_8UC1);
	vector<Mat> yuvChannels;
	
	/*NEED ASSERT RECT HERE*/
	Point2i origin = Point2i(center.x-CROP_W/2,center.y-CROP_H/2);
	assert(origin.x > 0 && origin.y > 0);

	Rect rCrop = Rect(origin.x,origin.y,CROP_W,CROP_H);
	Mat imgBGR_crop = imgBGR(rCrop); /*imshow("Color",imgBGR_crop);*/
	Mat imgD_crop = imgD(rCrop); /*imshow("Depth",imgD_crop);*/
	//Size realSize = Size(CROP_W/2,CROP_H);

	//Process Depth
	//resize(imgD_crop,imgD_crop,realSize,0,0,INTER_NEAREST);
	bool isDepthOk = genDepthMask(imgD_crop,depthThresh,depthMask);
	if (!isDepthOk) return;

	//Process Color
	cvtColor(imgBGR_crop,imgYUV,CV_BGR2YCrCb);
	split(imgYUV, yuvChannels);
	Mat uImg = yuvChannels[1];
	Mat vImg = yuvChannels[2];


	/*LOOP: step 1*/
	colorPointsMarking(uImg,vImg,depthMask,preview);
	//imshow("preview",preview);waitKey();

	/*LOOP: step 2*/
	colorNeighborMarking(uImg,vImg,depthMask,preview);
	//imshow("preview",preview);waitKey();

	/*LOOP: step 3*/
	//openning(preview,preview,1);
	eroding(preview,preview,1);
	dilating(preview,preview,2);
	//imshow("preview",preview);waitKey();
}

void SkinTracker::trackFramesYUV422(Mat & rawYUV, Mat & imgD, Point3f handPoint)
{
	/* YUV Crop Frame: 64x128 */
	int x_factor = int(handPoint.x+0.5f);
	int y_factor = int(handPoint.y+0.5f);
	if (x_factor > 639-64) x_factor = 639-64;
	if (x_factor < 64) x_factor = 64;
	if (y_factor > 479-64) y_factor = 479-64;
	if (y_factor < 64) y_factor = 64;
	
	ushort depthThresh = ushort(handPoint.z);
	//printf("depth = %d",depthThresh);
	trackFramesYUV422(rawYUV,imgD,Point2i(x_factor,y_factor),depthThresh);
}

void SkinTracker::trackFramesYUV422(Mat & rawYUV, Mat & imgD, Point2i center, ushort depthThresh)
{
	Mat preview = Mat::zeros(CROP_H,CROP_W/2,CV_8UC1);
	Mat depthMask = Mat::zeros(CROP_H,CROP_W/2,CV_8UC1);
	
	
	/*NEED ASSERT RECT HERE*/
	Point2i origin = Point2i(center.x-CROP_W/2,center.y-CROP_H/2);

	Rect depthCrop = Rect(origin.x,origin.y,CROP_W,CROP_H);
	Size uvSize = Size(CROP_W/2,CROP_H);
	Rect uvCrop = Rect(origin.x/2,origin.y,CROP_W/2,CROP_H);

	vector<Mat> yuvChannels;
	split(rawYUV, yuvChannels);
	Mat uImg = yuvChannels[2];
	Mat vImg = yuvChannels[0];
	//imshow("U",uImg);
	//imshow("V",vImg);

	Mat u_crop = uImg(uvCrop); /*imshow("U",u_crop);*/
	Mat v_crop = vImg(uvCrop); /*imshow("V",v_crop);*/
	Mat d_crop = imgD(depthCrop); 
	resize(d_crop,d_crop,uvSize,0,0,INTER_NEAREST);

	genDepthMask(d_crop,depthThresh,depthMask);
	bool isDepthOk = genDepthMask(d_crop,depthThresh,depthMask);
	if (!isDepthOk)	return;

	/*LOOP: step 1*/
	colorPointsMarking(u_crop,v_crop,depthMask,preview);

	/*LOOP: step 2*/
	colorNeighborMarking(u_crop,v_crop,depthMask,preview);

	/*LOOP: step 3*/
	//openning(preview,preview,1);
	eroding(preview,preview,1);
	dilating(preview,preview,2);
	//imshow("preview",preview);waitKey();
	resize(preview,preview,Size2i(128,128),0,0,INTER_NEAREST);
	previewMask = preview;
	d_crop_rect = depthCrop;
}

void SkinTracker::colorPointsMarking(Mat & uChannel, Mat & vChannel, Mat & mask, Mat & output)
{
	for (int i = 6; i < uChannel.rows-6; i++)
	{
		for (int j = 3; j < uChannel.cols-3; j++)
		{
			uchar dMask = mask.at<uchar>(i,j);
			if (dMask > 0)
			{
				uchar uColor = uChannel.at<uchar>(i,j);
				uchar vColor = vChannel.at<uchar>(i,j);
				double prob1 = m_Psc.at<double>(uColor,vColor);
				if (prob1>Tmax)	output.at<uchar>(i,j)= 255;
			}				
		}
	}
}

void SkinTracker::colorNeighborMarking(Mat & uChannel, Mat & vChannel, Mat & mask, Mat & outputMat)
{
	for (int i = 5; i < outputMat.rows-5; i++)
	{
		for (int j = 2; j < outputMat.cols-2; j++)
		{
			if (mask.at<uchar>(i,j) > 0 && outputMat.at<uchar>(i,j) != 255)
			{
				double prob = m_Psc.at<double>(uChannel.at<uchar>(i,j),vChannel.at<uchar>(i,j));
				if ( prob>Tmin && prob<Tmax)
				{
					/*check 4 direct points*/
					bool pass_4 = 
						outputMat.at<uchar>(i,j-1)== 255 || 
						outputMat.at<uchar>(i,j+1)== 255 ||
						outputMat.at<uchar>(i-1,j)== 255 ||	
						outputMat.at<uchar>(i+1,j)== 255;
					if (pass_4)
					{
						outputMat.at<uchar>(i,j)= 255;
					}
					else
					{
						/*check 4 more points*/
						bool pass_8 = 
							outputMat.at<uchar>(i-1,j-1)== 255 ||
							outputMat.at<uchar>(i-1,j+1)== 255 || 
							outputMat.at<uchar>(i+1,j-1)== 255 || 
							outputMat.at<uchar>(i+1,j+1)== 255;
						if (pass_8) outputMat.at<uchar>(i,j)= 255;
					}
				}
			}
		}

	}
}

/*SUPPORT*/
bool genDepthMask(Mat & inputMat, ushort depthThresh, Mat & outputMat)
{
	int count = 0;
	for (int i = 0; i < inputMat.rows; i++)
	{
		for (int j = 0; j < inputMat.cols; j++)
		{
			if (inputMat.at<ushort>(i,j) > depthThresh-80 && 
				inputMat.at<ushort>(i,j) < depthThresh+80)
			{
				outputMat.at<uchar>(i,j) = 255;
				count++;
			}	
		}
	}
	if (count<500) return false;
	dilating(outputMat,outputMat,5);
	return true;
}

void eroding(Mat & inputMat, Mat & outputMat, int erode_size)
{
	Mat element = getStructuringElement( MORPH_ELLIPSE,
		Size( erode_size+1, 2*erode_size+1 ),
		Point( erode_size, erode_size ) );
	erode( inputMat, outputMat, element );
}

void dilating(Mat & inputMat, Mat & outputMat, int dilation_size)
{
	Mat element = getStructuringElement( MORPH_ELLIPSE,
		Size( dilation_size+1, 2*dilation_size+1 ),
		Point( dilation_size, dilation_size ) );
	/// Apply the dilation operation
	dilate( inputMat, outputMat, element );
}

/*UTILITIES*/
void compareImage(Mat & colorImg, Mat & maskImg)
{
	//int count = 0;
	for (int i = 0; i < colorImg.rows; i++)
	{
		for (int j = 0; j < colorImg.cols; j++)
		{
			if (maskImg.at<uchar>(i,j) ==255)
			{
				colorImg.at<Vec3b>(i,j)[2] = 255;
			}
		}
	}
	//printf("count = %d; size = %dx%d\n",count,maskImg.cols,maskImg.rows);
}

void cannyTest(Mat & grayScaleImg)
{
	Mat canny_output;
	vector<vector<Point> > contours;

	vector<Vec4i> hierarchy;
	int thresh = 127;
	RNG rng(12345);
	// Detect edges using canny
	Canny( grayScaleImg, canny_output, thresh, thresh*2, 3 );
	// Find contours
	imshow( "outp", canny_output );
	findContours( canny_output, contours, hierarchy, CV_RETR_TREE, CV_CHAIN_APPROX_SIMPLE, Point(0, 0) );

	// Draw contours
	Mat drawing = Mat::zeros( canny_output.size(), CV_8UC3 );

	vector<vector<Point> > contoursFiltered;
	for( int i = 0; i< contours.size(); i++ )
	{
		if (contours[i].size()>15)
		{
			contoursFiltered.push_back(contours[i]);
		}
	}

	for (int i=0; i<contoursFiltered.size();i++)
	{
		Scalar color = Scalar( rng.uniform(0, 255), rng.uniform(0,255), rng.uniform(0,255) );
		drawContours( drawing, contoursFiltered, i, color, 1, 8, hierarchy, 0, Point() );
	}

	/// Show in a window
	namedWindow( "Contours", CV_WINDOW_AUTOSIZE );
	imshow( "Contours", drawing );
}

void previewColorImage(Mat & mRGB, COLOR_TYPE type)
{
	std::string windowsName;
	Mat imgConvert;
	vector<Mat> imgChannels;

	switch (type)
	{
	case COLOR_RGB:
		windowsName += "RGB";
		//cvtColor(mRGB,imgConvert,CV_BGR2YCrCb);
		split(imgConvert, imgChannels);
		break;
	case COLOR_YUV:
		windowsName += "YUV";
		cvtColor(mRGB,imgConvert,CV_BGR2YCrCb);
		split(imgConvert, imgChannels);
		break;
	case COLOR_HSV:
		windowsName += "HSV";
		cvtColor(mRGB,imgConvert,CV_BGR2HSV);
		split(imgConvert, imgChannels);
		break;
	default:
		return;
		break;
	}

	imshow(windowsName+"[1]",imgChannels[0]);
	imshow(windowsName+"[2]",imgChannels[1]);
	imshow(windowsName+"[3]",imgChannels[2]);
}

void previewHistogram(Mat & colorImage, COLOR_TYPE type)
{
	std::string windowsName = "Histogram ";
	Mat imgConvert;
	vector<Mat> imgChannels;

	int histSize = 256;
	float range[] = { 0, 256 } ;
	const float* histRange = { range };
	bool uniform = true; bool accumulate = false;

	Mat hist1, hist2, hist3;
	Scalar scalar1, scalar2, scalar3;
	switch (type)
	{
	case COLOR_RGB:
		windowsName+="RGB"; scalar1 = Scalar(255,0,0);scalar2 = Scalar(0,255,0);scalar3 = Scalar(0,0,255);break;
	case COLOR_YUV:
		cvtColor(colorImage,colorImage,CV_BGR2YCrCb);
		windowsName+="YUV"; scalar1 = Scalar(255,255,255);scalar2 = Scalar(0,255,0);scalar3 = Scalar(255,0,0);break;
	case COLOR_HSV:
		cvtColor(colorImage,colorImage,CV_BGR2HSV);
		windowsName+="HSV"; scalar1 = Scalar(255,255,255);scalar2 = Scalar(0,255,0);scalar3 = Scalar(255,0,0);break;
	default:
		windowsName+="RGB"; scalar1 = Scalar(255,0,0);scalar2 = Scalar(0,255,0);scalar3 = Scalar(0,0,255);break;
	}

	split(colorImage, imgChannels);

	/// Compute the histograms:
	calcHist( &imgChannels[0], 1, 0, Mat(), hist1, 1, &histSize, &histRange, uniform, accumulate );
	calcHist( &imgChannels[1], 1, 0, Mat(), hist2, 1, &histSize, &histRange, uniform, accumulate );
	calcHist( &imgChannels[2], 1, 0, Mat(), hist3, 1, &histSize, &histRange, uniform, accumulate );

	// Draw the histograms
	int hist_w = 512; int hist_h = 400;
	int bin_w = cvRound( (double) hist_w/histSize );

	Mat histImage( hist_h, hist_w, CV_8UC3, Scalar( 0,0,0) );

	/// Normalize the result to [ 0, histImage.rows ]
	normalize(hist1, hist1, 0, histImage.rows, NORM_MINMAX, -1, Mat() );
	normalize(hist2, hist2, 0, histImage.rows, NORM_MINMAX, -1, Mat() );
	normalize(hist3, hist3, 0, histImage.rows, NORM_MINMAX, -1, Mat() );

	/// Draw for each channel
	for( int i = 1; i < histSize; i++ )
	{
		line( histImage, 
			Point( bin_w*(i-1), hist_h - cvRound(hist1.at<float>(i-1)) ) ,
			Point( bin_w*(i), hist_h - cvRound(hist1.at<float>(i)) ),
			scalar1, 2, 8, 0  );
		line( histImage, 
			Point( bin_w*(i-1), hist_h - cvRound(hist2.at<float>(i-1)) ) ,
			Point( bin_w*(i), hist_h - cvRound(hist2.at<float>(i)) ),
			scalar2, 2, 8, 0  );
		line( histImage, 
			Point( bin_w*(i-1), hist_h - cvRound(hist3.at<float>(i-1)) ) ,
			Point( bin_w*(i), hist_h - cvRound(hist3.at<float>(i)) ),
			scalar3, 2, 8, 0  );
	}

	/// Display
	namedWindow(windowsName, CV_WINDOW_NORMAL );
	imshow(windowsName, histImage );
}